import torch
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)
from peft import LoraConfig, get_peft_model
import json


# 读取 JSON 文件
with open("/root/autodl-tmp/data/data/train.json", "r", encoding="utf-8") as f:
    data = json.load(f)

# 数据处理函数 - 修复重复的<|im_end|>问题
def format_conversation(conversation):
    formatted_text = ""
    for turn in conversation:
        if turn["role"] == "user":
            formatted_text += f"<|im_start|>user\n{turn['content']}<|im_end|>\n"
        elif turn["role"] == "assistant":
            # 确保assistant内容中不包含<|im_end|>
            content = turn['content'].replace('<|im_end|>', '')
            formatted_text += f"<|im_start|>assistant\n{content}<|im_end|>\n"
    return formatted_text

# 加载模型和分词器
model_name = "/root/autodl-tmp/data/models/models/Qwen/Qwen3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# 确保有pad_token
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype=torch.float16,  # 使用dtype而不是torch_dtype
    device_map="auto",
    trust_remote_code=True
)

# 配置LoRA - 平衡配置
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_dropout=0.1,  # 适中的dropout
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# 准备训练数据
train_texts = []
for conversation in data:
    formatted_text = format_conversation(conversation)
    train_texts.append(formatted_text)

train_dataset = Dataset.from_dict({"text": train_texts})

print("训练数据示例:")
for i, text in enumerate(train_texts[:2]):
    print(f"示例 {i+1}:")
    print(text)
    print()

# 分词函数 - 使用动态padding
def tokenize_function(examples):
    # 使用动态padding，确保batch内长度一致
    tokenized = tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,
        padding=False,  # 不使用静态padding
        return_tensors=None
    )
    tokenized["labels"] = tokenized["input_ids"].copy()
    return tokenized

tokenized_dataset = train_dataset.map(
    tokenize_function,
    batched=False,  # 改为False，避免batch内padding问题
    remove_columns=train_dataset.column_names
)

# 数据收集器 - 使用动态padding
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=8,  # 确保长度是8的倍数
)

# 训练参数 - 平衡参数，避免过拟合但保持语义相似
training_args = TrainingArguments(
    output_dir="/root/autodl-tmp/data/save_modles",
    overwrite_output_dir=True,
    per_device_train_batch_size=1,  # 使用较小的batch_size避免内存问题
    gradient_accumulation_steps=8,
    num_train_epochs=15,  # 适中的训练轮数
    learning_rate=2e-4,   # 适中的学习率
    fp16=True,
    logging_steps=5,
    save_steps=100,
    save_total_limit=2,
    remove_unused_columns=False,
    dataloader_pin_memory=False,
    warmup_ratio=0.1,
    weight_decay=0.01,    # 添加权重衰减防止过拟合
)

# 创建Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

print("开始训练...")
trainer.train()

# 保存模型
trainer.save_model()
tokenizer.save_pretrained("/root/autodl-tmp/data/save_modles")
print("训练完成！模型已保存到 ./qwen2-payment-balanced")